# 1. 导入核心依赖
from typing import Annotated
from typing_extensions import TypedDict  # Python 3.11+ 可改用 from typing import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages  # 用于追加对话消息
from langchain.chat_models import init_chat_model  # 初始化 LLM
from langchain.schema import AIMessage  # 识别 AI 消息类型
from dotenv import load_dotenv
import os


# 3. 定义 State（状态）：存储对话历史
class ChatState(TypedDict):
    # messages：对话历史列表，更新方式为 add_messages（追加，不覆盖）
    messages: Annotated[list, add_messages]


# 4. 初始化 LLM（OpenAI GPT-4o-mini）
llm = init_chat_model(
    model="deepseek-chat",  # 模型名称
    model_provider="deepseek",  # 模型提供商
    # API 密钥（从环境变量读取）
    temperature=0.7  # 随机性：0~1，越低越严谨
)


# 5. 定义 Node（节点）：调用 LLM 生成回答
def chatbot_node(state: ChatState):
    """
    输入：当前 State（包含历史对话 messages）
    输出：需要更新的 State 字段（新的 AI 消息）
    """
    # 1. 从 State 中读取历史对话
    history_messages = state["messages"]

    # 2. 调用 LLM 生成回答（传入历史对话，支持多轮上下文）
    ai_response = llm.invoke(history_messages)

    # 3. 返回需要更新的字段：将 AI 回答追加到 messages 中
    return {"messages": [ai_response]}


# 6. 构建 Graph（图）：组装节点和边
# 6.1 初始化图（绑定 State 结构）
graph_builder = StateGraph(ChatState)

# 6.2 添加节点（名称："chatbot"，逻辑：chatbot_node 函数）
graph_builder.add_node("chatbot", chatbot_node)

# 6.3 添加边（定义流程走向）
graph_builder.add_edge(START, "chatbot")  # 从 START 到 chatbot 节点
graph_builder.add_edge("chatbot", END)  # 从 chatbot 节点到 END

# 6.4 编译图（生成可执行的工作流）
chat_graph = graph_builder.compile()

# 7. 可视化图结构（可选，方便调试）
try:
    from IPython.display import Image, display

    # 生成 Mermaid 可视化图片
    graph_image = chat_graph.get_graph().draw_mermaid_png()
    # 保存图片到文件
    with open("graph.png", "wb") as f:
        f.write(graph_image)
    print("图表已保存为 graph.png 文件，请在项目根目录查看")
    # 如果在 Jupyter 环境中，也可以显示图片
    try:
        display(Image(data=graph_image))
    except:
        pass
    with open("graph.png", "wb") as f:
        f.write(graph_image)
except Exception as e:
    print(f"可视化失败（非必要，可忽略）：{e}")


# 8. 定义交互逻辑：多轮对话
def run_chatbot():
    print("=== 简易 LangGraph 聊天机器人 ===")
    print("提示：输入空字符串即可退出")

    while True:
        # 8.1 获取用户输入
        user_input = input("\nUser: ")
        if user_input.strip() == "":  # 输入为空则退出
            print("对话结束！")
            break

        # 8.2 调用 Graph 处理输入（流式输出，体验更好）
        print("Assistant: ", end="")  # 不换行，逐字输出回答
        for event in chat_graph.stream(
                # 初始 State：传入用户输入的消息
                {"messages": [{"role": "user", "content": user_input}]}
        ):
            # 遍历流式事件，提取 AI 回答
            for node_output in event.values():
                if "messages" in node_output:
                    # 获取最新的消息（AI 回答）
                    latest_msg = node_output["messages"][-1]
                    if isinstance(latest_msg, AIMessage):
                        # 逐字输出（模拟流式效果）
                        print(latest_msg.content, end="")
    print("\n")


# 9. 启动聊天机器人
if __name__ == "__main__":
    run_chatbot()
